import typing
from torch.onnx import symbolic_helper
from torch.onnx.symbolic_helper import parse_args
import torch
_OPSET_VERSION = 12
_registered_ops: typing.AbstractSet[str] = set()

def register():
    """
    Register ONNX Runtime's built-in contrib ops.
	Should be run before torch.onnx.export().
	"""
    # @staticmethod
    # @parse_args('v', 'i', 'i', 'i', 'b')
    def grid_sampler(g, input, grid, mode, padding_mode, align_corners):
		# mode
		#   'bilinear'      : onnx::Constant[value={0}]
		#   'nearest'       : onnx::Constant[value={1}]
		#   'bicubic'       : onnx::Constant[value={2}]
		# padding_mode
		#   'zeros'         : onnx::Constant[value={0}]
		#   'border'        : onnx::Constant[value={1}]
		#   'reflection'    : onnx::Constant[value={2}]
        # grid = symbolic_helper._maybe_get_const(grid, "t")
        mode = symbolic_helper._maybe_get_const(mode, "i")
        padding_mode = symbolic_helper._maybe_get_const(padding_mode, "i")
        # mode_str = ["bilinear", "nearest", "bicubic"][mode]
        # padding_mode_str = ["zeros", "border", "reflection"][padding_mode]
        align_corners = int(symbolic_helper._maybe_get_const(align_corners, "b"))

		# From opset v13 onward, the output shape can be specified with
		# (N, C, H, W) (N, H_out, W_out, 2) => (N, C, H_out, W_out)
		# input_shape = input.type().sizes()
		# gird_shape = grid.type().sizes()
		# output_shape = input_shape[:2] + gird_shape[1:3]
		# g.op(...).setType(input.type().with_sizes(output_shape))
        inputs = [input, grid]
        kwargs = {
            "mode_i":mode, 
            "padding_mode_i":padding_mode, 
            "align_corners_i": align_corners}

        return g.op(
		    ## op name, modify here. not sure whether "com.microsoft::" is required
			"nvinfer1::GridSamplePluginDynamic",  
			*inputs, **kwargs
		)

    _reg(grid_sampler)

def _reg(symbolic_fn: typing.Callable):
	name = "::%s" % symbolic_fn.__name__
	torch.onnx.register_custom_op_symbolic(name, symbolic_fn, _OPSET_VERSION)
	_registered_ops.add(name)